public class AVLTree {
    static class TreeNode{
        public int val;
        public int bf;
        public TreeNode left;
        public TreeNode right;
        public TreeNode parent;

        public TreeNode(int val){
            this.val = val;
        }
    }

    public TreeNode root;

    public boolean insert(int val){
        // 1. 插入节点
        TreeNode node = new TreeNode(val);
        if(root == null){
            root = node;
            return true;
        }

        TreeNode parent = null;
        TreeNode cur = root;

        while(cur != null){
            if(cur.val < val){
                parent = cur;
                cur = cur.right;
            }else if(cur.val > val){
                parent = cur;
                cur = cur.left;
            }else{
                return false;
            }
        }

        // cur == null
        if(parent.val < val){
            parent.right = node;
        }else{
            parent.left = node;
        }
        node.parent = parent;
        cur = node;

        // 2. 平衡因子的修改
        while(parent != null){
            if(cur == parent.left){
                parent.bf--;
            }else{
                parent.bf++;
            }

            if(parent.bf == 0){
                return true;
            }else if(parent.bf == 1 || parent.bf == -1){
                cur = parent;
                parent = parent.parent;
            }else{
                if(parent.bf == 2){
                    if(cur.bf == 1){
                        // 左单旋
                        rotateLeft(parent);
                    }else{
                        // cur.bf == -1
                        rotateRL(parent);
                    }
                }else{ // parent.bf == -2
                    if(cur.bf == -1){
                        // 右单旋
                        rotateRight(parent);
                    }else{
                        rotateLR(parent);
                    }
                }
                return true;
            }
        }

        return true;
    }

    private void rotateRight(TreeNode parent){
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;

        // 修改四个指向
        parent.left = subLR;
        subL.right = parent;

        if(subLR != null){
            subLR.parent = parent;
        }

        // 记录父节点的父节点，旋转完成会用到
        TreeNode pParent = parent.parent;

        parent.parent = subL;

        // 修改 subL 的指向
        if(root == parent){
            root = subL;
            root.parent = null;
        }else{
            if(pParent.left == parent){
                pParent.left = subL;
            }else{
                pParent.right = subL;
            }
        }
        subL.parent = pParent;
        // 修改平衡因子
        subL.bf = 0;
        parent.bf = 0;
    }

    private void rotateLeft(TreeNode parent){
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;

        // 修改四个指向
        parent.right = subRL;
        subR.left = parent;

        if(subRL != null){
            subRL.parent = parent;
        }

        // 修改 parent 的父节点之前，保存 parent 的父节点
        TreeNode pParent = parent.parent;

        parent.parent = subR;

        if(root == parent){
            root = subR;
            root.parent = null;
        }else{
            if(pParent.left == parent){
                pParent.left = subR;
            }else{
                pParent.right = subR;
            }
        }

        subR.parent = pParent;

        subR.bf = 0;
        parent.bf = 0;
    }

    private void rotateLR(TreeNode parent){
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;
        int bf = subLR.bf;

        rotateLeft(subL);
        rotateRight(parent);

        if(bf == -1){
            parent.bf = 1;
            subL.bf = 0;
            subLR.bf = 0;
        }else if(bf == 1){
            subL.bf = -1;
            subLR.bf = 0;
            parent.bf = 0;
        }
    }

    private void rotateRL(TreeNode parent){
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;
        int bf = subRL.bf;

        rotateRight(subR);
        rotateLeft(parent);

        if(bf == -1){
            parent.bf = 0;
            subRL.bf = 0;
            subR.bf = 1;
        }else if(bf == 1){
            parent.bf = -1;
            subRL.bf = 0;
            subR.bf = 0;
        }
    }

    // 中序遍历 - 有序表示当前树是二叉搜索树
    public void inOrder(TreeNode root){
        if(root == null){
            return;
        }
        inOrder(root.left);
        System.out.print(root.val + " ");
        inOrder(root.right);
    }

    public int getHeight(TreeNode root){
        if(root == null) return 0;

        int left = getHeight(root.left);
        if(left == -1) return -1;

        int right = getHeight(root.right);
        if(right == -1) return -1;

        if(Math.abs(left - right) > 1) return -1;

        return Math.max(left, right) + 1;
    }
}
